import asyncio
from aiortc import RTCPeerConnection, RTCSessionDescription
from aiortc.contrib.signaling import TcpSocketSignaling

class WebRTCHandler:
    def __init__(self, on_datachannel, on_video_request):
        self.pc = RTCPeerConnection()
        self.on_datachannel = on_datachannel
        self.on_video_request = on_video_request
        self.signaling = TcpSocketSignaling("0.0.0.0", 8080)
        
        # 设置事件处理器
        @self.pc.on("datachannel")
        def pc_on_datachannel(channel):
            self.on_datachannel(channel)
        
        @self.pc.on("track")
        def pc_on_track(track):
            print("Track received:", track.kind)
    
    async def start_signaling(self):
        await self.signaling.connect()
        
        @self.signaling.on("offer")
        async def on_offer(offer):
            await self.pc.setRemoteDescription(offer)
            
            # 如果请求中包含视频轨道，启动视频流
            for description in offer.sdp.splitlines():
                if "m=video" in description:
                    self.on_video_request()
            
            # 创建应答
            answer = await self.pc.createAnswer()
            await self.pc.setLocalDescription(answer)
            await self.signaling.send(self.pc.localDescription)
        
        @self.signaling.on("candidate")
        async def on_candidate(candidate):
            await self.pc.addIceCandidate(candidate)
        
        # 主循环
        while True:
            obj = await self.signaling.receive()
            if isinstance(obj, RTCSessionDescription):
                if obj.type == "offer":
                    await on_offer(obj)
            elif isinstance(obj, RTCIceCandidate):
                await on_candidate(obj)